#!/usr/bin/env python
import rospy
import cv2 as cv
import numpy as np
import os
import yaml
from pyquaternion import Quaternion
from cv_bridge import CvBridge, CvBridgeError
from prometheus_msgs.msg import DetectionInfo
from sensor_msgs.msg import Image


detect_info = DetectionInfo()
detect_info.object_name = "landpat"

aruco_dict = cv.aruco.getPredefinedDictionary(cv.aruco.DICT_6X6_250)
aruco_param = cv.aruco.DetectorParameters()
aruco_detector = cv.aruco.ArucoDetector(aruco_dict, aruco_param)

def image_cb(msg, bridge, detect_pub, _config_data):
    try:
        cv_img = bridge.imgmsg_to_cv2(msg, "bgr8")
        shape = cv_img.shape
        landpat_length = _config_data['landpat_length']
        camera_matrix = _config_data['camera_matrix']
        distortion_coefficients = _config_data['distortion_coefficients']

        (corners, ids, rejected) = aruco_detector.detectMarkers(cv_img)
        # print(ids, corners)
        if ids is not None:
            corner = corners[0]
            _id = ids[0]
            cv.aruco.drawDetectedMarkers(cv_img, corners, ids)
            rvecs, tvecs, _ = cv.aruco.estimatePoseSingleMarkers(corner, landpat_length, camera_matrix, distortion_coefficients)
            cv.drawFrameAxes(cv_img, camera_matrix, distortion_coefficients, rvecs, tvecs, 0.1)
            detect_info.header.stamp = rospy.Time.now()
            detect_info.detected = True
            
            detect_info.position[0] = rvecs[0][0][0]
            detect_info.position[1] = rvecs[0][0][1]
            detect_info.position[2] = rvecs[0][0][2]
            # print(rvecs)
        else:
            detect_info.header.stamp = rospy.Time.now()
            detect_info.detected = False
        
        detect_pub.publish(detect_info)
        cv.imshow('cv_img', cv_img)
        key = cv.waitKey(1)

    except CvBridgeError as e:
        rospy.logerr(e)

if __name__ == '__main__':
    rospy.init_node('detect_fire_point')
    rate = rospy.Rate(20)
    bridge = CvBridge()

    camera_config_file = rospy.get_param('~camera_config_file', '/home/xs/catkin_ws/src/demo1/config/camera_param_real.yaml')
    uav_id = rospy.get_param('~uav_id', 1)
    rospy.loginfo("Camera config file path is: " + camera_config_file)

    if not os.path.exists(camera_config_file):
        rospy.logerr("Can not find camera config file {0}".format(camera_config_file))
        exit(0)

    with open(camera_config_file, 'r', encoding='utf-8') as fr:
        config_data = yaml.load(fr, yaml.FullLoader)

    landpat_length = float(config_data['aruco_det_len'])
    fx = float(config_data['fx'])
    cx = float(config_data['x0'])
    fy = float(config_data['fy'])
    cy = float(config_data['y0'])

    k1 = float(config_data['k1'])
    k2 = float(config_data['k2'])
    p1 = float(config_data['p1'])
    p2 = float(config_data['p2'])
    k3 = float(config_data['k3'])

    camera_matrix = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=float)
    distortion_coefficients = np.array([k1, k2, p1, p2, k3], dtype=float)

    rospy.loginfo('landpat length is : {0} [m]'.format(landpat_length))
    rospy.loginfo('camera matrix is \n{0}, \ndistortion coefficients is \n{1}'.format(camera_matrix, distortion_coefficients))

    input_param = {'landpat_length': landpat_length, 'camera_matrix': camera_matrix, 'distortion_coefficients': distortion_coefficients}

    detect_pub = rospy.Publisher('/uav'+ str(uav_id) + '/prometheus/object_detection/aruco_landpat', DetectionInfo, queue_size=2)

    img_sub = rospy.Subscriber('/iris1/camera/image_raw', Image, lambda msg: image_cb(msg, bridge, detect_pub, input_param))

    rospy.spin()
   
